import pandas as pd


# 导入所需的数据文件
train_data = pd.read_csv('train.csv', dtype={'sale_date': str, 'goodsid': int, 'is_pro': int})
test_data = pd.read_csv('test.csv', dtype={'sale_date': str, 'goodsid': int, 'is_pro': int})
goods_info = pd.read_csv('goods_ch.csv',
                         usecols=['goodsid', 'div_id', 'catg_l_id',
                                  'catg_m_id', 'catg_s_id', 'season_class'])
# 是否为周末字段变为0-1数据，节假日字段均只取后面的数字编号，非节假日为0
date_ch = pd.read_csv('date_ch.csv', 
                      encoding='gbk', 
                      usecols=['dim_date_id', 'year_code', 'quanter_code', 'month_code', 
                               'day_week_cn', 'week_day_code', 'is_weekend', 'official_holiday_code',
                               'festival_code'], 
                      converters={'dim_date_id': lambda x: pd.to_datetime(str(x), format='%Y%m%d'),
                                  'is_weekend': lambda x: 1 if x == '是' else 0,
                                  'official_holiday_code': lambda x: int(x[1:]) if x else 0,
                                  'festival_code': lambda x: int(x[1:]) if x else 0})
train_data['sale_date'] = pd.to_datetime(train_data['sale_date'], format='%Y%m%d')
test_data['sale_date'] = pd.to_datetime(test_data['sale_date'], format='%Y%m%d')
pp_train = pd.read_csv('prophet_predict_train.csv', 
                       converters={'ds': pd.to_datetime, 'goodsid': lambda x: int(float(x))})
pp_test = pd.read_csv('prophet_predict_test.csv', 
                      converters={'ds': pd.to_datetime, 'goodsid': lambda x: int(float(x))})

# 训练集宽表生成
train_wide_data = pd.merge(train_data, goods_info, on='goodsid', how='left')
train_wide_data = pd.merge(train_wide_data, date_ch, left_on='sale_date', right_on='dim_date_id', how='left')\
                  .drop('dim_date_id', axis=1)
train_wide_data['day'] = train_wide_data['sale_date'].apply(lambda x: x.day)
train_wide_data = pd.merge(train_wide_data, pp_train, left_on=['sale_date', 'goodsid'], right_on=['ds', 'goodsid'], how='left')\
                  .drop('ds', axis=1)
train_wide_data = train_wide_data.drop('sale_date', axis=1)
train_wide_data.to_csv('train_wide_data.csv', index=False)

# 测试集宽表生成
test_wide_data = pd.merge(test_data, goods_info, on='goodsid', how='left')
test_wide_data = pd.merge(test_wide_data, date_ch, left_on='sale_date', right_on='dim_date_id', how='left')\
                  .drop('dim_date_id', axis=1)
test_wide_data['day'] = test_wide_data['sale_date'].apply(lambda x: x.day)
test_wide_data = pd.merge(test_wide_data, pp_test, left_on=['sale_date', 'goodsid'], right_on=['ds', 'goodsid'], how='left')\
                  .drop('ds', axis=1)
test_wide_data = test_wide_data.drop('sale_date', axis=1)
test_wide_data.to_csv('test_wide_data_t.csv', index=False)
